import pandas as pd
import lancedb
from langchain_community.vectorstores import LanceDB
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.documents import Document


def vectorstore_to_csv():
    print("🚀 VectorStore 转 CSV 导出...")

    # 1. 连接到现有 VectorStore 或创建示例
    db = lancedb.connect("./data/lancedb")
    embeddings = HuggingFaceEmbeddings(
        model_name="D:\\models\\models\\sentence-transformers\\paraphrase-multilingual-MiniLM-L12-v2"
    )

    try:
        # 尝试打开现有 VectorStore
        vector_store = LanceDB.from_existing_index(
            embeddings,
            connection=db,
            table_name="ai_documents"
        )
        print("✅ 成功连接到现有 VectorStore")
    except:
        print("⚠️  创建示例 VectorStore")
        vector_store = create_sample_vectorstore(db, embeddings)

    # 2. 通过搜索获取所有文档（近似方法）
    # 注意：这不是获取所有文档的标准方法，但对于演示目的可行
    print("🔄 检索文档数据...")

    # 使用一个通用查询来获取大部分文档
    all_docs = []
    generic_queries = ["学习", "技术", "编程", "数据", "智能"]

    for query in generic_queries:
        docs = vector_store.similarity_search(query, k=10)
        all_docs.extend(docs)

    # 去重
    unique_docs = []
    seen_contents = set()
    for doc in all_docs:
        if doc.page_content not in seen_contents:
            seen_contents.add(doc.page_content)
            unique_docs.append(doc)

    print(f"📄 检索到 {len(unique_docs)} 个唯一文档")

    # 3. 转换为 DataFrame
    data = []
    for doc in unique_docs:
        row = {
            'content': doc.page_content,
            'metadata': str(doc.metadata),
            'content_preview': doc.page_content[:100] + '...' if len(doc.page_content) > 100 else doc.page_content
        }
        # 展开 metadata
        for key, value in doc.metadata.items():
            row[f'metadata_{key}'] = value
        data.append(row)

    df = pd.DataFrame(data)

    # 4. 导出到 CSV
    output_file = "vectorstore_export.csv"
    df.to_csv(output_file, index=False, encoding='utf-8')
    print(f"✅ VectorStore 数据已导出到: {output_file}")

    # 显示预览
    print("\n📋 导出数据预览:")
    print(df[['content_preview', 'metadata']].head())

    return df


def create_sample_vectorstore(db, embeddings):
    """创建示例 VectorStore"""
    documents = [
        Document(
            page_content="机器学习算法可以从数据中自动学习模式和规律",
            metadata={"category": "AI", "difficulty": "advanced", "source": "book"}
        ),
        Document(
            page_content="Python编程语言简单易学，适合数据分析和人工智能开发",
            metadata={"category": "Programming", "difficulty": "beginner", "source": "tutorial"}
        ),
        Document(
            page_content="深度学习在图像识别和自然语言处理中取得了重大突破",
            metadata={"category": "AI", "difficulty": "advanced", "source": "research"}
        ),
        Document(
            page_content="数据科学需要统计学、编程和领域知识的综合能力",
            metadata={"category": "Data Science", "difficulty": "intermediate", "source": "course"}
        )
    ]

    vector_store = LanceDB.from_documents(
        documents=documents,
        embedding=embeddings,
        connection=db,
        table_name="sample_documents"
    )

    return vector_store

# 运行 VectorStore 导出
df = vectorstore_to_csv()